#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <malloc.h>
#include <execinfo.h>
#include <elf.h>


// Symbol table
typedef struct {
    char *name;
    void *addr;
} qmp_func_t;

typedef struct {
    size_t size;
    qmp_func_t **func;
} qmp_symtab_t;


static qmp_symtab_t *qmp_symtab;


static int qmp_func_cmp(const qmp_func_t **f1, const qmp_func_t **f2)
{
    return (*f1)->addr < (*f2)->addr ? -1 : (*f1)->addr == (*f2)->addr ? 0 : 1;
}

static void qmp_elf_init()
{
    FILE *fp;
    char *strtab_buf, *binary_name;
    int   i, sym_elnum;
    Elf32_Ehdr elf_header;
    Elf32_Shdr elf_symtab, elf_strtab;

    binary_name = getenv("QMP_EXE");
    if (!binary_name) {
        printf("cannot read envvar QMP_EXE\n");
        exit(EXIT_FAILURE);
    }
    fp = fopen(binary_name, "rb");
    if (!fp) {
        printf("cannot open %s for reading\n", binary_name);
        exit(EXIT_FAILURE);
    }

    fread(&elf_header, sizeof(elf_header), 1, fp);
    fseek(fp, elf_header.e_shoff, SEEK_SET);

    elf_symtab.sh_type = elf_strtab.sh_type = SHT_NULL;

    for (i = 0; i < elf_header.e_shnum; i++) {
        Elf32_Shdr elf_section;

        fread(&elf_section, sizeof(elf_section), 1, fp);

        if (elf_section.sh_type == SHT_SYMTAB)
            elf_symtab = elf_section;
        else if (elf_section.sh_type == SHT_STRTAB)
            elf_strtab = elf_section;
    }
    if (elf_symtab.sh_type != SHT_SYMTAB) {
        printf("no symbol table\n");
        exit(EXIT_FAILURE);
    }
    if (elf_strtab.sh_type != SHT_STRTAB) {
        printf("no string table\n");
        exit(EXIT_FAILURE);
    }

    fseek(fp, elf_strtab.sh_offset, SEEK_SET);
    strtab_buf = malloc(elf_strtab.sh_size);
    fread(strtab_buf, elf_strtab.sh_size, 1, fp);

    sym_elnum = elf_symtab.sh_size / elf_symtab.sh_entsize;

    qmp_symtab = (qmp_symtab_t *)malloc(sizeof(qmp_symtab_t));
    qmp_symtab->size = 0;
    qmp_symtab->func = (qmp_func_t **)calloc(sizeof(qmp_func_t *), sym_elnum);

    fseek(fp, elf_symtab.sh_offset, SEEK_SET);
    for (i = 0; i < sym_elnum; i++) {
        Elf32_Sym elf_sym;

        fread(&elf_sym, sizeof(elf_sym), 1, fp);
        if (ELF32_ST_TYPE(elf_sym.st_info) == STT_FUNC && elf_sym.st_value) {
            qmp_symtab->func[qmp_symtab->size] = (qmp_func_t *)malloc(sizeof(qmp_func_t));
            qmp_symtab->func[qmp_symtab->size]->name = strdup(strtab_buf + elf_sym.st_name);
            qmp_symtab->func[qmp_symtab->size++]->addr = (void *)elf_sym.st_value;
        }
    }

    qsort(qmp_symtab->func, qmp_symtab->size, sizeof(qmp_func_t *),
          (int(*)(const void *, const void *))qmp_func_cmp);

    free(strtab_buf);
    fclose(fp);
}

static qmp_func_t *qmp_symtab_get(const void *addr)
{
    size_t i;
    qmp_func_t *func = NULL;

    if (addr < qmp_symtab->func[0]->addr ||
        addr > qmp_symtab->func[qmp_symtab->size - 1]->addr + 512)
        return NULL;

    for (i = 0; !func && i < qmp_symtab->size; i++) {
        if (qmp_symtab->func[i]->addr > addr)
            func = qmp_symtab->func[i - 1];
    }

    return func;
}

// Prime numbers
static unsigned int qmp_prime_num[] = {
    11, 19, 37, 73, 109, 163, 251, 367, 557, 823, 1237, 1861, 2777, 4177, 6247,
    9371, 14057, 21089, 31627, 47431, 71143, 106721, 160073, 240101, 360163,
    540217, 810343, 1215497, 1823231, 2734867, 4102283, 6153409, 9230113,
    13845163
};

// Trace
typedef struct {
    size_t max;
    size_t cur;
    size_t sum;
    size_t fc_size;
    qmp_func_t **fc_buf;
} qmp_trace_t;

typedef struct {
    size_t size;
    size_t pnum;
    qmp_trace_t **data;
} qmp_trace_map_t;


static qmp_trace_map_t *qmp_trace_map;


#define QMP_TRACEMAP_CAP (qmp_prime_num[qmp_trace_map->pnum])

static void qmp_trace_map_init()
{
    qmp_trace_map = (qmp_trace_map_t *)malloc(sizeof(qmp_trace_map_t));
    qmp_trace_map->size = 0;
    qmp_trace_map->pnum = 0;
    qmp_trace_map->data = (qmp_trace_t **)calloc(QMP_TRACEMAP_CAP * sizeof(qmp_trace_t *), 1);
}

static qmp_trace_t *qmp_trace_create(void **buf, size_t size)
{
    size_t i;
    qmp_trace_t *trace = (qmp_trace_t *)calloc(sizeof(qmp_trace_t), 1);

    trace->fc_buf = (qmp_func_t **)calloc(sizeof(qmp_func_t *), size);

    for (i = 0; i < size; i++) {
        qmp_func_t *func = qmp_symtab_get(buf[i]);

        if (func)
            trace->fc_buf[trace->fc_size++] = func;
    }

    return trace;
}

#define HASH_MIX_IN(hash, key) \
    do { \
        (hash) += (key); \
        (hash) += (hash) << 10; \
        (hash) ^= (hash) >> 6; \
    } while (0)

#define HASH_SHAKE(hash) \
    do { \
        (hash) += (hash) << 3; \
        (hash) ^= (hash) >> 11; \
        (hash) += (hash) << 15; \
    } while (0)

static unsigned int qmp_trace_map_hash(qmp_trace_t *trace)
{
    int i;
    unsigned int h = 0;

    for (i = 0; i < trace->fc_size; i++)
        HASH_MIX_IN(h, (unsigned int)trace->fc_buf[i]->addr);

    HASH_SHAKE(h);
    return h;
}

static int qmp_trace_map_eq(qmp_trace_t *t1, qmp_trace_t *t2)
{
    int i;

    if (t1->fc_size != t2->fc_size)
        return 0;

    for (i = 0; i < t1->fc_size; i++) {
        if (t1->fc_buf[i]->addr != t2->fc_buf[i]->addr)
            return 0;
    }
    return 1;
}

static qmp_trace_t **qmp_trace_map_lookup(qmp_trace_t *retrace)
{
    unsigned int  idx   = qmp_trace_map_hash(retrace) % QMP_TRACEMAP_CAP;
    qmp_trace_t **trace = &qmp_trace_map->data[idx];

    if (!*trace)
        return trace;

    while (*trace) {
        if (idx + 1 < QMP_TRACEMAP_CAP)
            idx++;
        else
            idx = 0;
        trace = &qmp_trace_map->data[idx];
    }

    return trace;
}

static qmp_trace_t *qmp_trace_map_add(void **buf, size_t size)
{
    qmp_trace_t  *retrace = qmp_trace_create(buf, size);
    qmp_trace_t **trace   = qmp_trace_map_lookup(retrace);

    qmp_trace_map->size++;
    *trace = retrace;


    if (qmp_trace_map->size * 3 > QMP_TRACEMAP_CAP) {
        int i, old_data_size   = qmp_prime_num[qmp_trace_map->pnum];
        qmp_trace_t **old_data = qmp_trace_map->data;

        qmp_trace_map->pnum++;
        qmp_trace_map->data =
            (qmp_trace_t **)calloc(QMP_TRACEMAP_CAP * sizeof(qmp_trace_t *), 1);

        for (i = 0; i < old_data_size; i++) {
            qmp_trace_t **trace, *old_trace = old_data[i];

            if (!old_trace)
                continue;

            trace  = qmp_trace_map_lookup(old_trace);
            *trace = old_trace;
        }
        free(old_data);
    }

    return retrace;
}

static qmp_trace_t *qmp_trace_map_get(void **buf, size_t size)
{
    qmp_trace_t *retrace = qmp_trace_create(buf, size);
    unsigned int idx     = qmp_trace_map_hash(retrace) % QMP_TRACEMAP_CAP;
    qmp_trace_t *trace   = qmp_trace_map->data[idx];

    if (!trace) {
        free(retrace);
        return NULL;
    }

    while (trace && !qmp_trace_map_eq(trace, retrace)) {
        if (idx + 1 < QMP_TRACEMAP_CAP)
            idx++;
        else
            idx = 0;
        trace = qmp_trace_map->data[idx];
    }

    free(retrace);
    return trace;
}

// Ptr
typedef struct {
    size_t   size;
    qmp_trace_t *trace;
} qmp_ptr_t;

typedef struct qmp_ptr_pair_t qmp_ptr_pair_t;
struct qmp_ptr_pair_t {
    void           *p;
    qmp_ptr_t      *ptr;
    qmp_ptr_pair_t *next;
};

typedef struct {
    size_t size;
    size_t pnum;
    qmp_ptr_pair_t **data;
} qmp_ptr2qmp_trace_t;


static qmp_ptr2qmp_trace_t *qmp_ptr2trace;


#define QMP_PTR2TRACE_CAP (qmp_prime_num[qmp_ptr2trace->pnum])

static void qmp_ptr2trace_init()
{
    qmp_ptr2trace = (qmp_ptr2qmp_trace_t *)malloc(sizeof(qmp_ptr2qmp_trace_t));
    qmp_ptr2trace->size = 0;
    qmp_ptr2trace->pnum = 0;
    qmp_ptr2trace->data = (qmp_ptr_pair_t **)calloc(QMP_PTR2TRACE_CAP * sizeof(qmp_ptr_pair_t*), 1);
}

static void qmp_ptr2trace_add(void *p, size_t size, qmp_trace_t *trace)
{
    unsigned int idx = (unsigned int)p % QMP_PTR2TRACE_CAP;
    qmp_ptr_t      *ptr  = (qmp_ptr_t *)malloc(sizeof(qmp_ptr_t));
    qmp_ptr_pair_t *pair = (qmp_ptr_pair_t *)malloc(sizeof(qmp_ptr_pair_t));

    ptr->size  = size;
    ptr->trace = trace;

    pair->p    = p;
    pair->ptr  = ptr;

    pair->next = qmp_ptr2trace->data[idx];
    qmp_ptr2trace->data[idx] = pair;
    qmp_ptr2trace->size++;


    if (qmp_ptr2trace->size * 3 > QMP_PTR2TRACE_CAP) {
        int i, old_data_size = qmp_prime_num[qmp_ptr2trace->pnum];
        qmp_ptr_pair_t **old_data = qmp_ptr2trace->data;

        qmp_ptr2trace->pnum++;
        qmp_ptr2trace->data =
            (qmp_ptr_pair_t **)calloc(QMP_PTR2TRACE_CAP * sizeof(qmp_ptr_pair_t *), 1);

        for (i = 0; i < old_data_size; i++) {
            pair = old_data[i];
            if (!pair)
                continue;

            while (pair) {
                qmp_ptr_pair_t *npair = pair->next;

                idx = (unsigned int)(pair->p) % QMP_PTR2TRACE_CAP;

                pair->next = qmp_ptr2trace->data[idx];
                qmp_ptr2trace->data[idx] = pair;

                pair = npair;
            }
        }
        free(old_data);
    }
}

static qmp_ptr_t *qmp_ptr2trace_get(void *p)
{
    unsigned int idx = (unsigned int)p % QMP_PTR2TRACE_CAP;
    qmp_ptr_pair_t *pair = qmp_ptr2trace->data[idx];

    while (pair && pair->p != p)
        pair = pair->next;

    return pair ? pair->ptr : NULL;
}

static void qmp_ptr2trace_remove(void *p)
{
    unsigned int idx = (unsigned int)p % QMP_PTR2TRACE_CAP;
    qmp_ptr_pair_t *pair, *ppair = NULL, **fpair = &qmp_ptr2trace->data[idx];

    pair = *fpair;
    while (pair && pair->p != p) {
        ppair = pair;
        pair  = pair->next;
    }

    if (!pair)
        return ;

    qmp_ptr2trace->size--;

    if (!ppair)
        *fpair = pair->next;
    else 
        ppair->next = pair->next;

    free(pair->ptr);
    free(pair);
}

// Hooks
static void *(*old_malloc_hook)(size_t, const void *);
static void  (*old_free_hook)(void *, const void *);

static void *qmp_malloc_hook(size_t, const void *);
static void  qmp_free_hook(void *, const void *);


static void qmp_hook()
{
    __malloc_hook = qmp_malloc_hook;
    __free_hook   = qmp_free_hook;
}

static void qmp_unhook()
{
    __malloc_hook = old_malloc_hook;
    __free_hook   = old_free_hook;
}

#define BTSIZE 100

static void *qmp_malloc_hook(size_t size, const void *caller)
{
    void *p;
    void *buf[BTSIZE];
    int   b_size;
    qmp_trace_t *trace = NULL;

    qmp_unhook();

    b_size = backtrace(buf, BTSIZE);
    if (b_size > 0) {
        trace = qmp_trace_map_get(buf, b_size);
        if (!trace)
            trace = qmp_trace_map_add(buf, b_size);
    }

    p = malloc(size);

    if (trace && p) {
        trace->cur += size;
        trace->sum += size;
        trace->max = trace->max < trace->cur ? trace->cur : trace->max;
        qmp_ptr2trace_add(p, size, trace);
    }

    qmp_hook();

    return p;
}

static void qmp_free_hook(void *p, const void *caller)
{
    qmp_ptr_t *ptr;

    qmp_unhook();
    
    ptr = qmp_ptr2trace_get(p);
    if (ptr) {
        ptr->trace->cur -= ptr->size;
        qmp_ptr2trace_remove(p);
    }

    free(p);
    
    qmp_hook();
}

// Init

static int qmp_trace_cmp(const qmp_trace_t **t1, const qmp_trace_t **t2)
{
    return (*t1)->max < (*t2)->max ? -1 : (*t1)->max == (*t2)->max ? 0 : 1;
}

static void qmp_print_result()
{
    int i, k, l;
    qmp_trace_t **trace_buf = (qmp_trace_t **)calloc(sizeof(qmp_trace_t *), QMP_TRACEMAP_CAP);

    for (i = 0, k = 0; i < QMP_TRACEMAP_CAP; i++) {
        qmp_trace_t *trace = qmp_trace_map->data[i];

        if (!trace)
            continue;

        trace_buf[k++] = qmp_trace_map->data[i];
    }

    qsort(trace_buf, k, sizeof(qmp_trace_t *), (int(*)(const void *, const void *))qmp_trace_cmp);

    for (i = 0; i < k; i++) {
        qmp_trace_t *trace = trace_buf[i];
        size_t      max = trace->max, sum = trace->sum;
        const char *max_ss = "",   *sum_ss = "";

        if (trace->max / 1024 > 1024) {
            max = trace->max / 1024;
            max_ss = "Kb";
        }
        if (trace->sum / 1024 > 1024) {
            sum = trace->sum / 1024;
            sum_ss = "Kb";
        }

        printf("%u%s %u%s | ", max, max_ss, sum, sum_ss);
        for (l = 0; l < trace->fc_size; l++)
            printf("::%s", trace->fc_buf[l]->name);
        printf("\n");
    }
}

static void qmp_init()
{
    qmp_elf_init();
    qmp_ptr2trace_init();
    qmp_trace_map_init();
    atexit(qmp_print_result);

    old_malloc_hook = __malloc_hook;
    old_free_hook   = __free_hook;

    qmp_hook();
}

// Initialize point
void (*__malloc_initialize_hook) () = qmp_init;

